"""Session Key Manager.

Handles encryption key generation, exchange, rotation and secure storage
for BLE communication sessions.
"""

import asyncio
import hashlib
import hmac
import json
import logging
import secrets
import time
from typing import Dict, List, Optional, Tuple, Any
from enum import Enum
from dataclasses import dataclass, field
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives.asymmetric import ec
from cryptography.hazmat.backends import default_backend

logger = logging.getLogger(__name__)


# Key Management Configuration
KEY_CONFIG = {
    "algorithm": "ECDH-P256",        # Key exchange algorithm
    "session_key_size": 16,          # Session key size (128 bits)
    "rotation_interval": 3600,       # Key rotation interval (1 hour)
    "max_key_age": 7200,            # Maximum key age (2 hours)
    "derivation_rounds": 10000,      # PBKDF2 rounds
    "salt_size": 16,                # Salt size in bytes
    "iv_size": 12,                  # GCM IV size (96 bits)
    "tag_size": 16,                 # GCM authentication tag size
    "max_usage_count": 1000000      # Maximum key usage count
}


class KeyState(Enum):
    """Key states."""
    GENERATING = 0        # Key being generated
    ACTIVE = 1           # Key is active
    ROTATING = 2         # Key rotation in progress
    EXPIRED = 3          # Key has expired
    REVOKED = 4          # Key has been revoked


class CryptoAlgorithm(Enum):
    """Supported crypto algorithms."""
    AES_128_GCM = "AES-128-GCM"
    AES_256_GCM = "AES-256-GCM"
    CHACHA20_POLY1305 = "ChaCha20-Poly1305"


@dataclass
class SessionKey:
    """Session encryption key."""
    key_id: str
    key_data: bytes
    algorithm: CryptoAlgorithm
    created_time: float
    expires_at: float
    state: KeyState = KeyState.ACTIVE
    usage_count: int = 0
    max_usage: int = field(default_factory=lambda: KEY_CONFIG["max_usage_count"])
    
    @property
    def is_expired(self) -> bool:
        """Check if key is expired."""
        return (time.time() > self.expires_at or 
                self.usage_count >= self.max_usage or
                self.state == KeyState.EXPIRED)
    
    @property
    def age(self) -> float:
        """Get key age in seconds."""
        return time.time() - self.created_time
    
    def increment_usage(self) -> bool:
        """Increment usage counter."""
        self.usage_count += 1
        if self.usage_count >= self.max_usage:
            self.state = KeyState.EXPIRED
            return False
        return True


@dataclass
class KeyExchangeData:
    """Key exchange information."""
    exchange_id: str
    session_id: str
    public_key: bytes
    private_key: ec.EllipticCurvePrivateKey
    algorithm: str
    created_time: float
    expires_at: float
    
    @property
    def is_expired(self) -> bool:
        """Check if exchange data is expired."""
        return time.time() > self.expires_at


class KeyRotationPolicy:
    """Key rotation policy."""
    
    def __init__(self, 
                 time_based: bool = True,
                 usage_based: bool = True,
                 max_age: int = None,
                 max_usage: int = None):
        self.time_based = time_based
        self.usage_based = usage_based
        self.max_age = max_age or KEY_CONFIG["max_key_age"]
        self.max_usage = max_usage or KEY_CONFIG["max_usage_count"]
    
    def should_rotate(self, key: SessionKey) -> bool:
        """Check if key should be rotated."""
        if self.time_based and key.age >= self.max_age:
            return True
        if self.usage_based and key.usage_count >= self.max_usage * 0.9:
            return True
        return False


class SessionKeyManager:
    """Session Key Manager for secure BLE communications."""
    
    def __init__(self, algorithm: CryptoAlgorithm = CryptoAlgorithm.AES_128_GCM):
        self.algorithm = algorithm
        self.session_keys: Dict[str, SessionKey] = {}
        self.key_exchanges: Dict[str, KeyExchangeData] = {}
        self.rotation_policy = KeyRotationPolicy()
        
        # Background tasks
        self.rotation_task: Optional[asyncio.Task] = None
        self.is_running = False
        
        logger.info(f"SessionKeyManager initialized with {algorithm.value}")
    
    async def start(self):
        """Start key manager."""
        if self.is_running:
            return
            
        self.is_running = True
        self.rotation_task = asyncio.create_task(self._rotation_loop())
        
        logger.info("SessionKeyManager started")
    
    async def stop(self):
        """Stop key manager."""
        self.is_running = False
        
        if self.rotation_task:
            self.rotation_task.cancel()
            try:
                await self.rotation_task
            except asyncio.CancelledError:
                pass
        
        logger.info("SessionKeyManager stopped")
    
    def generate_key_pair(self) -> Tuple[ec.EllipticCurvePrivateKey, bytes]:
        """Generate ECDH key pair."""
        # Generate private key
        private_key = ec.generate_private_key(ec.SECP256R1(), default_backend())
        
        # Get public key bytes
        public_key = private_key.public_key()
        public_bytes = public_key.public_numbers().x.to_bytes(32, 'big')
        public_bytes += public_key.public_numbers().y.to_bytes(32, 'big')
        
        return private_key, public_bytes
    
    def derive_shared_secret(self, private_key: ec.EllipticCurvePrivateKey,
                           peer_public_key: bytes) -> bytes:
        """Derive shared secret from key exchange."""
        try:
            # Reconstruct peer public key
            x = int.from_bytes(peer_public_key[:32], 'big')
            y = int.from_bytes(peer_public_key[32:64], 'big')
            
            peer_public_numbers = ec.EllipticCurvePublicNumbers(x, y, ec.SECP256R1())
            peer_public_key_obj = peer_public_numbers.public_key(default_backend())
            
            # Perform ECDH
            shared_key = private_key.exchange(ec.ECDH(), peer_public_key_obj)
            return shared_key
            
        except Exception as e:
            logger.error(f"Key derivation error: {e}")
            raise ValueError("Invalid public key for exchange")
    
    def derive_session_key(self, shared_secret: bytes, 
                          session_id: str, salt: bytes = None) -> bytes:
        """Derive session key from shared secret."""
        if salt is None:
            salt = secrets.token_bytes(KEY_CONFIG["salt_size"])
        
        # Use PBKDF2 for key derivation
        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=KEY_CONFIG["session_key_size"],
            salt=salt + session_id.encode('utf-8'),
            iterations=KEY_CONFIG["derivation_rounds"],
            backend=default_backend()
        )
        
        return kdf.derive(shared_secret)
    
    async def create_key_exchange(self, session_id: str) -> KeyExchangeData:
        """Create key exchange for session."""
        # Generate key pair
        private_key, public_key = self.generate_key_pair()
        
        exchange_id = secrets.token_urlsafe(16)
        exchange_data = KeyExchangeData(
            exchange_id=exchange_id,
            session_id=session_id,
            public_key=public_key,
            private_key=private_key,
            algorithm=KEY_CONFIG["algorithm"],
            created_time=time.time(),
            expires_at=time.time() + 300  # 5 minute expiry
        )
        
        self.key_exchanges[exchange_id] = exchange_data
        
        logger.info(f"Key exchange created: {exchange_id} for session {session_id}")
        return exchange_data
    
    async def complete_key_exchange(self, exchange_id: str, 
                                  peer_public_key: bytes) -> SessionKey:
        """Complete key exchange and generate session key."""
        exchange_data = self.key_exchanges.get(exchange_id)
        if not exchange_data:
            raise ValueError(f"Key exchange not found: {exchange_id}")
        
        if exchange_data.is_expired:
            del self.key_exchanges[exchange_id]
            raise ValueError(f"Key exchange expired: {exchange_id}")
        
        try:
            # Derive shared secret
            shared_secret = self.derive_shared_secret(
                exchange_data.private_key, peer_public_key
            )
            
            # Derive session key
            session_key_data = self.derive_session_key(
                shared_secret, exchange_data.session_id
            )
            
            # Create session key
            key_id = secrets.token_urlsafe(16)
            session_key = SessionKey(
                key_id=key_id,
                key_data=session_key_data,
                algorithm=self.algorithm,
                created_time=time.time(),
                expires_at=time.time() + KEY_CONFIG["rotation_interval"]
            )
            
            # Store session key
            self.session_keys[exchange_data.session_id] = session_key
            
            # Clean up exchange data
            del self.key_exchanges[exchange_id]
            
            logger.info(f"Key exchange completed: {exchange_id} -> {key_id}")
            return session_key
            
        except Exception as e:
            logger.error(f"Key exchange completion error: {e}")
            del self.key_exchanges[exchange_id]
            raise
    
    async def generate_session_key(self, session_id: str, 
                                 password: str = None) -> SessionKey:
        """Generate session key directly (for testing/fallback)."""
        if password:
            # Derive key from password
            salt = secrets.token_bytes(KEY_CONFIG["salt_size"])
            key_data = self.derive_session_key(
                password.encode('utf-8'), session_id, salt
            )
        else:
            # Generate random key
            key_data = secrets.token_bytes(KEY_CONFIG["session_key_size"])
        
        key_id = secrets.token_urlsafe(16)
        session_key = SessionKey(
            key_id=key_id,
            key_data=key_data,
            algorithm=self.algorithm,
            created_time=time.time(),
            expires_at=time.time() + KEY_CONFIG["rotation_interval"]
        )
        
        self.session_keys[session_id] = session_key
        
        logger.info(f"Session key generated: {key_id} for session {session_id}")
        return session_key
    
    def get_session_key(self, session_id: str) -> Optional[SessionKey]:
        """Get active session key."""
        key = self.session_keys.get(session_id)
        if key and not key.is_expired:
            return key
        return None
    
    async def rotate_session_key(self, session_id: str) -> SessionKey:
        """Rotate session key."""
        old_key = self.session_keys.get(session_id)
        if old_key:
            old_key.state = KeyState.ROTATING
        
        # Generate new key
        new_key = await self.generate_session_key(session_id)
        
        if old_key:
            old_key.state = KeyState.EXPIRED
            logger.info(f"Key rotated: {old_key.key_id} -> {new_key.key_id}")
        
        return new_key
    
    def revoke_session_key(self, session_id: str) -> bool:
        """Revoke session key."""
        key = self.session_keys.get(session_id)
        if key:
            key.state = KeyState.REVOKED
            del self.session_keys[session_id]
            logger.info(f"Session key revoked: {key.key_id}")
            return True
        return False
    
    def encrypt_data(self, session_id: str, plaintext: bytes) -> Optional[Dict[str, Any]]:
        """Encrypt data with session key."""
        key = self.get_session_key(session_id)
        if not key:
            return None
        
        try:
            # Generate random IV
            iv = secrets.token_bytes(KEY_CONFIG["iv_size"])
            
            # Create cipher
            cipher = Cipher(
                algorithms.AES(key.key_data),
                modes.GCM(iv),
                backend=default_backend()
            )
            encryptor = cipher.encryptor()
            
            # Encrypt data
            ciphertext = encryptor.update(plaintext) + encryptor.finalize()
            
            # Increment key usage
            key.increment_usage()
            
            return {
                "algorithm": key.algorithm.value,
                "iv": iv.hex(),
                "ciphertext": ciphertext.hex(),
                "tag": encryptor.tag.hex(),
                "key_id": key.key_id
            }
            
        except Exception as e:
            logger.error(f"Encryption error: {e}")
            return None
    
    def decrypt_data(self, session_id: str, encrypted_data: Dict[str, Any]) -> Optional[bytes]:
        """Decrypt data with session key."""
        key = self.get_session_key(session_id)
        if not key:
            return None
        
        try:
            # Extract components
            iv = bytes.fromhex(encrypted_data["iv"])
            ciphertext = bytes.fromhex(encrypted_data["ciphertext"])
            tag = bytes.fromhex(encrypted_data["tag"])
            
            # Create cipher
            cipher = Cipher(
                algorithms.AES(key.key_data),
                modes.GCM(iv, tag),
                backend=default_backend()
            )
            decryptor = cipher.decryptor()
            
            # Decrypt data
            plaintext = decryptor.update(ciphertext) + decryptor.finalize()
            
            # Increment key usage
            key.increment_usage()
            
            return plaintext
            
        except Exception as e:
            logger.error(f"Decryption error: {e}")
            return None
    
    async def _rotation_loop(self):
        """Background key rotation loop."""
        while self.is_running:
            try:
                await asyncio.sleep(60)  # Check every minute
                
                # Find keys that need rotation
                keys_to_rotate = []
                for session_id, key in self.session_keys.items():
                    if self.rotation_policy.should_rotate(key):
                        keys_to_rotate.append(session_id)
                
                # Rotate keys
                for session_id in keys_to_rotate:
                    try:
                        await self.rotate_session_key(session_id)
                    except Exception as e:
                        logger.error(f"Key rotation error for {session_id}: {e}")
                
                # Clean up expired exchanges
                expired_exchanges = [eid for eid, ex in self.key_exchanges.items() 
                                   if ex.is_expired]
                for eid in expired_exchanges:
                    del self.key_exchanges[eid]
                
                if keys_to_rotate or expired_exchanges:
                    logger.debug(f"Rotated {len(keys_to_rotate)} keys, "
                               f"cleaned {len(expired_exchanges)} exchanges")
                
            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.error(f"Rotation loop error: {e}")
    
    def get_statistics(self) -> Dict[str, Any]:
        """Get key manager statistics."""
        active_keys = sum(1 for k in self.session_keys.values() 
                         if k.state == KeyState.ACTIVE)
        expired_keys = sum(1 for k in self.session_keys.values() 
                          if k.state == KeyState.EXPIRED)
        
        return {
            "key_manager": {
                "is_running": self.is_running,
                "algorithm": self.algorithm.value,
                "total_keys": len(self.session_keys),
                "active_keys": active_keys,
                "expired_keys": expired_keys,
                "pending_exchanges": len(self.key_exchanges)
            },
            "key_usage": {
                "total_usage": sum(k.usage_count for k in self.session_keys.values()),
                "average_age": sum(k.age for k in self.session_keys.values()) / 
                              max(1, len(self.session_keys))
            }
        }
    
    async def cleanup(self):
        """Clean up key manager."""
        await self.stop()
        
        # Clear all keys securely
        for key in self.session_keys.values():
            # Overwrite key data
            key.key_data = b'\x00' * len(key.key_data)
        
        self.session_keys.clear()
        self.key_exchanges.clear()
        
        logger.info("SessionKeyManager cleanup completed")


# Example usage and testing
if __name__ == "__main__":
    import asyncio
    
    async def test_key_manager():
        """Test key manager functionality."""
        key_manager = SessionKeyManager()
        
        # Start key manager
        await key_manager.start()
        
        try:
            session_id = "test_session_001"
            
            # Test key exchange flow
            print("=== Key Exchange Test ===")
            
            # Server side: create exchange
            exchange_data = await key_manager.create_key_exchange(session_id)
            print(f"Exchange created: {exchange_data.exchange_id}")
            print(f"Public key length: {len(exchange_data.public_key)} bytes")
            
            # Client side: generate key pair (simulated)
            client_private, client_public = key_manager.generate_key_pair()
            print(f"Client public key generated: {len(client_public)} bytes")
            
            # Server side: complete exchange
            session_key = await key_manager.complete_key_exchange(
                exchange_data.exchange_id, client_public
            )
            print(f"Session key created: {session_key.key_id}")
            print(f"Key algorithm: {session_key.algorithm.value}")
            
            # Test encryption/decryption
            print("\n=== Encryption Test ===")
            
            test_message = b"Hello, secure BLE world!"
            print(f"Original message: {test_message}")
            
            # Encrypt
            encrypted = key_manager.encrypt_data(session_id, test_message)
            if encrypted:
                print(f"Encrypted successfully")
                print(f"IV: {encrypted['iv'][:16]}...")
                print(f"Ciphertext length: {len(encrypted['ciphertext'])} chars")
                
                # Decrypt
                decrypted = key_manager.decrypt_data(session_id, encrypted)
                if decrypted:
                    print(f"Decrypted: {decrypted}")
                    print(f"Match: {decrypted == test_message}")
                else:
                    print("Decryption failed")
            else:
                print("Encryption failed")
            
            # Test key rotation
            print("\n=== Key Rotation Test ===")
            old_key_id = session_key.key_id
            new_key = await key_manager.rotate_session_key(session_id)
            print(f"Key rotated: {old_key_id} -> {new_key.key_id}")
            
            # Verify new key works
            encrypted_new = key_manager.encrypt_data(session_id, test_message)
            if encrypted_new:
                decrypted_new = key_manager.decrypt_data(session_id, encrypted_new)
                print(f"New key works: {decrypted_new == test_message}")
            
            # Print statistics
            stats = key_manager.get_statistics()
            print(f"\nStatistics: {json.dumps(stats, indent=2)}")
            
        finally:
            # Cleanup
            await key_manager.cleanup()
    
    # Run test
    asyncio.run(test_key_manager())